import requests, json, base64
from python.utils import image_change

TOKEN = "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"

API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
# API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
# API_URL = "https://api-inference.huggingface.co/models/williamberman/stable-diffusion-v1-5-controlnet-canny-full-no-ema"
headers = {"Authorization": "Bearer " + TOKEN}  # 自己的密钥


def data_package(content):
    data = {
        "prompt": content,
        "negative_prompt": None,
        "width": "512",
        "height": "512",
        "samples": "1",
        "num_inference_steps": "20",
        "safety_checker": "no",
        "enhance_prompt": "yes",
        "seed": None,
        "guidance_scale": 7.5,
        "multi_lingual": "no",
        "panorama": "no",
        "self_attention": "no",
        "upscale": "no",
        "embeddings_model": None,
        "webhook": None,
        "track_id": None
    }
    return json.dumps(data)


def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content


def get_result(image_binary):
    return image_change.binary_to_base64(image_binary)


def image_generate(content):
    data = data_package(content)
    print("!!!!")
    image_binary = query(data)
    print(image_binary)
    return get_result(image_binary)
